from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.tools import tool
from utils import websearch
from typing import Annotated
from langgraph.prebuilt import create_react_agent


@tool
def web_search(query: Annotated[str, "互联网查询内容"]):
    """通过web_search工具查询互联网上的信息"""
    _rt = websearch(query)
    return _rt


_researcher_system_template = """
你是一个优秀的信息搜集者（你只管收集信息，不要做其他工作）,
你的职责是根据用户的任务需求和团队的要求搜集相关的信息，
如果你完成信息收集过后，就将任务交给其他角色。
如果任务全部顺利完成，返回的字符串中给出消息 ANSWER FINNAL,这样团队就知道该停下来了。
你可以使用以下的工具协助你完成任务:
{tools_name}

"""


class Researcher:

    def __init__(self, llm):

        _tools = [
            web_search,
        ]

        _prompt = ChatPromptTemplate.from_messages(
            [
                ("system", _researcher_system_template),
                MessagesPlaceholder(variable_name="messages"),
            ]
        )

        _prompt = _prompt.partial(tools_name=",".join([_tool.name for _tool in _tools]))
        _llm_with_tools_agent = create_react_agent(llm, tools=_tools)
        self._chain = _prompt | _llm_with_tools_agent

    def __call__(self, state):
        _rt = self._chain.invoke(state)
        _messages = _rt["messages"]
        return _messages[-1]


if __name__ == "__main__":
    from langchain_openai import ChatOpenAI

    _llm = ChatOpenAI(
        api_key="ollama",
        model="qwen2.5:7b",
        base_url="http://192.168.10.11:60026/v1",
        temperature=0.7,
    )

    _researcher = Researcher(_llm)
    _rt = _researcher(
        {
            "messages": [
                (
                    "human",
                    "获取英国过去5年的国内生产总值。一旦你把它编码好，并执行画图，就完成。",
                )
            ]
        }
    )

    print(_rt)
